'''
https://leetcode.cn/problems/paint-fence
'''

from functools import cache


class Solution:
    def numWays(self, n: int, k: int) -> int:
        @cache
        def total_ways(i):
            if i == 1:
                return k
            if i == 2:
                return k * k
            return (k - 1) * (total_ways(i - 1) + total_ways(i - 2))

        return total_ways(n)

    def numWays2(self, n: int, k: int) -> int:
        if n == 1: return k
        dp = [0] * (n + 1)
        dp[1] = k
        dp[2] = k * k
        for i in range(3, n + 1):
            dp[i] = (k - 1) * (dp[i - 1] + dp[i - 2])
        return dp[n]


n = 3
k = 2
print(Solution().numWays(n, k))